from ml_level_matcher import MLLevelMatcher
import json
import pandas as pd
from sklearn.metrics import classification_report, confusion_matrix
import joblib
import numpy as np
from typing import Dict, Any, List

def simplify_features(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """简化特征，只保留最重要的特征"""
    simplified_data = []
    
    for item in data:
        simplified_item = {
            'target_level': item['target_level'],
            'experience_score': item['experience_score'],
            'education_score': item['education_score'],
            'compensation_score': item['compensation_score'],
            'responsibility_score': item['responsibility_score'],
            'location_tier': item['location_tier'],
            'tech_level': item['tech_level'],
            'role_type': item['role_type'],
            'business_domain': '其他'  # 使用默认值
        }
        
        # 设置固定值的特征
        simplified_item['technical_score'] = 0
        simplified_item['workload_score'] = 70
        
        simplified_data.append(simplified_item)
        
    return simplified_data

def train_and_evaluate():
    """训练并评估简化版模型"""
    print("\n" + "="*50)
    print("开始训练简化版模型")
    print("="*50)
    
    print("\n加载数据...")
    with open('ali_data.json', 'r', encoding='utf-8') as f:
        training_data = json.load(f)
    
    print("简化特征...")
    simplified_data = simplify_features(training_data)
    
    # 创建模型实例
    matcher = MLLevelMatcher()
    
    print("\n开始训练模型...")
    matcher.train(simplified_data)
    
    print("保存模型...")
    matcher.save_model('model_simple.joblib')
    
    print("\n模型评估:")
    
    # 准备评估数据
    y_true = []
    y_pred = []
    
    for item in simplified_data:
        features = {k: v for k, v in item.items() if k != 'target_level'}
        predicted_level, confidence, _ = matcher.predict_level(features)
        y_true.append(item['target_level'])
        y_pred.append(predicted_level)
    
    print("\n分类报告:")
    print(classification_report(y_true, y_pred))
    
    cm = confusion_matrix(y_true, y_pred)
    unique_levels = sorted(list(set(y_true)))
    
    print("\n混淆矩阵:")
    print("预测值")
    print("实际值")
    print(" "*8 + " ".join(f"{level:>8}" for level in unique_levels))
    for i, level in enumerate(unique_levels):
        print(f"{level:>8}" + "".join(f"{cm[i][j]:>8}" for j in range(len(unique_levels))))
    
    print("\n各职级准确率:")
    for i, level in enumerate(unique_levels):
        level_total = sum(1 for y in y_true if y == level)
        level_correct = sum(1 for y, yp in zip(y_true, y_pred) if y == level and y == yp)
        accuracy = level_correct / level_total if level_total > 0 else 0
        print(f"{level}: {accuracy:.2%} ({level_correct}/{level_total})")
    
    importance = matcher.get_feature_importance()
    print("\n特征重要性:")
    for feature, score in sorted(importance.items(), key=lambda x: x[1], reverse=True):
        print(f"{feature}: {score:.4f}")
    
    print("\n测试预测示例:")
    test_features = {
        "experience_score": 75,      
        "education_score": 80,       
        "technical_score": 0,          
        "workload_score": 70,          
        "compensation_score": 65,      
        "responsibility_score": 80,    
        
        "location_tier": "一线城市",
        "tech_level": "高级",
        "role_type": "IC",
        "business_domain": "其他"     
    }
    
    predicted_level, confidence, class_probs = matcher.predict_level(test_features)
    print(f"\n测试预测结果:")
    print(f"预测职级: {predicted_level}")
    print(f"置信度: {confidence:.2f}")
    print("\n各职级概率:")
    for level, prob in sorted(class_probs.items()):
        print(f"{level}: {prob:.2f}")

if __name__ == "__main__":
    train_and_evaluate()